package com.wan.admin;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.lang.StringUtils;
import org.apache.flume.Event;
import org.apache.flume.event.SimpleEvent;
import org.apache.flume.instrumentation.SourceCounter;
import org.apache.rocketmq.client.consumer.DefaultMQPullConsumer;
import org.apache.rocketmq.client.consumer.PullResult;
import org.apache.rocketmq.client.consumer.PullStatus;
import org.apache.rocketmq.client.exception.MQBrokerException;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;
import org.apache.rocketmq.remoting.exception.RemotingException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.UnsupportedEncodingException;
import java.text.SimpleDateFormat;
import java.util.*;

/**
 * @项目名称：RocketMQ-Flume-Source
 * @类名称：TestUtil
 * @类描述：测试工具
 * @创建人：万星明
 * @创建时间：2020/8/19
 */
public class TestUtil {

    private static final Logger LOG = LoggerFactory.getLogger(TestUtil.class);

    private static String topic="LOGS_MESSAGE";

    private static String tags="*";

    private static String topicHeaderName="LOGS_MESSAGE";

    private static String tagsHeaderName="*";

    private static String nameserver="10.0.31.116:9876;10.0.31.117:9876";

    private static String consumerGroup="MyConsumerGroup_12";

    private static String messageModel="BROADCASTING";


    private static Integer batchSize =32;

    private static String p1 ="appId";

    private static String p2 ="createTime";

    private static Boolean isPrintLog = true;

    /**
     * 数据源数量
     */
    private static SourceCounter sourceCounter;

    private static DefaultMQPullConsumer consumer;

    /**
     * 获取默认消费者
     * @return 消费者
     */
    public static DefaultMQPullConsumer getConsumer() {
        //使用GroupName来初始化默认拉取消费者
        DefaultMQPullConsumer consumer = new DefaultMQPullConsumer(consumerGroup);
        //设置地址,消息订阅方式
        consumer.setNamesrvAddr(nameserver);
        consumer.setMessageModel( MessageModel.valueOf(messageModel));
        System.out.println("消费者设置完毕,服务器地址:"+nameserver+",消费组:"+consumerGroup);
        //返回消费者
        return consumer;
    }


    /**
     * 获取消息队列消费进度
     * @param queue 消息队列
     * @return 进度
     * @throws MQClientException MQ客户端异常
     */
    private static long getMessageQueueOffset(MessageQueue queue) throws MQClientException {
        // 从Broker获取Offset
        long offset = consumer.fetchConsumeOffset(queue, false);
        if (offset < 0L) {
            offset = 0L;
        }
        return offset;
    }


    /**
     * 存储拉取进度
     * @param queue 消息队列
     * @param offset 进度
     * @throws MQClientException 消息
     */
    private static void putMessageQueueOffset(MessageQueue queue,long offset) throws MQClientException {
        //存储Offset,客户端每隔5s会定时刷新到Broker或者写入本地缓存文件
        consumer.updateConsumeOffset(queue, offset);
    }


    public static void doConfigure() {
        consumer = Preconditions.checkNotNull(getConsumer());
        //注册监听队列
        consumer.registerMessageQueueListener(topic, null);
        if (isPrintLog){
            System.out.println("消费者成功注册:"+ JSON.toJSONString(consumer));
        }

        try {
            // 启动Consumer
            consumer.start();
        } catch (MQClientException e) {
            if (isPrintLog){
                System.out.println("启动消费者失败!");
            }
            Throwables.propagate(e);
        }
    }


    public static void main(String[] args) throws MQClientException, RemotingException, InterruptedException, MQBrokerException, UnsupportedEncodingException {
        doConfigure();

        //新建事件列表
        List<Event> eventList = Lists.newArrayList();
        Map<MessageQueue, Long> offsetMap = Maps.newHashMap();
        Event event;
        Map<String, String> headers;

        if (isPrintLog){
            System.out.println("读取数据处理程序开始!");
        }
        try
        {
            //从指定topic中拉取所有消息队列
            Set<MessageQueue> queues = Preconditions.checkNotNull(consumer.fetchSubscribeMessageQueues(topic));
            if (isPrintLog){
                System.out.println("从指定主题中拉取所有的消息队列对象:"+JSON.toJSONString(queues));
            }
            for (MessageQueue queue : queues) {
                //获取offset进度
                long offset = getMessageQueueOffset( queue );
                //消费者拉取结果
                PullResult pullResult = consumer.pullBlockIfNotFound( queue,tags,offset,batchSize );
                if (isPrintLog){
                    System.out.println("迭代队列————>消费者拉取结果:"+JSON.toJSONString(pullResult));
                }

                if (pullResult.getPullStatus() == PullStatus.FOUND) {
                    if (isPrintLog){
                        System.out.println("迭代队列————>拉取状态为FOUND");
                    }
                    //取出拉取结果消息列表
                    for (MessageExt messageExt : pullResult.getMsgFoundList()) {
                        //新建简单事件
                        event = new SimpleEvent();
                        headers = new HashMap<String,String>();
                        //将消息的主题和标签存入头部
                        headers.put( topicHeaderName,messageExt.getTopic() );
                        headers.put( tagsHeaderName,messageExt.getTags() );

                        //获取事件内容文本
                        String body = new String( messageExt.getBody(),"UTF-8" );
                        if (LOG.isDebugEnabled()) {
                            LOG.debug( "MessageQueue={}, Topic={}, Tags={}, Message: {}",new Object[] {
                                    queue,messageExt.getTopic(),messageExt.getTags(),body} );
                        }
                        if (isPrintLog){
                            System.out.println("迭代队列————>在迭代的消息数据————>MessageQueue:"+JSON.toJSONString(queue)+";Topic:"+JSON.toJSONString(messageExt.getTopic())
                                    +";Tags:"+JSON.toJSONString(messageExt.getTags())+";Message:"+JSON.toJSONString(body));
                        }

                        //设置赋值自定义配置
                        String appId = null;
                        String createTime = null;
                        try {
                            JSONObject jb = JSONObject.parseObject( body );
                            SimpleDateFormat format = new SimpleDateFormat( "yyyyMMdd" );
                            appId = jb.getString( p1 );
                            createTime = format.format( jb.getLong( p2 ) );
                        }
                        catch (Exception e) {
                            LOG.error( "JSON解析异常,Message:{}",e.getMessage(),e );
                            if (isPrintLog){
                                System.out.println("迭代队列————>JSON解析异常————>Message:"+JSON.toJSONString(e));
                            }
                        }

                        if (StringUtils.isBlank( appId )) {
                            appId = "unknown";
                        }
                        if (StringUtils.isBlank( createTime )) {
                            SimpleDateFormat format = new SimpleDateFormat( "yyyyMMdd" );
                            createTime = format.format( new Date() );
                        }
                        headers.put( p1,appId );
                        headers.put( p2,createTime );
                        //将消息内容存入事件内容
                        event.setBody( messageExt.getBody() );
                        event.setHeaders( headers );

                        if (isPrintLog){
                            System.out.println("迭代队列————>将消息存入事件中:"+JSON.toJSONString(event));
                        }

                        eventList.add( event );
                        if (isPrintLog){
                            System.out.println("迭代队列————>事件列表:"+JSON.toJSONString(eventList));
                        }
                    }
                    //进度集合存入
                    offsetMap.put( queue,pullResult.getNextBeginOffset() );
                    if (isPrintLog){
                        System.out.println("迭代队列————>存储进度:"+JSON.toJSONString(offsetMap));
                    }
                }
            }

            //如果事件列表不为空
            if (eventList.size()>0) {
                if (isPrintLog){
                    System.out.println("事件列表不为空处理!");
                }
                sourceCounter.incrementAppendBatchReceivedCount();
                sourceCounter.addToEventReceivedCount(eventList.size());
//                getChannelProcessor().processEventBatch(eventList);
                sourceCounter.incrementAppendBatchAcceptedCount();
                sourceCounter.addToEventAcceptedCount(eventList.size());
                eventList.clear();
                if (isPrintLog){
                    System.out.println("事件列表不为空处理结束:"+JSON.toJSONString(eventList));
                }
            }

            for (Map.Entry<MessageQueue, Long> entry : offsetMap.entrySet()) {
                //更新offset
                putMessageQueueOffset(entry.getKey(), entry.getValue());
            }
        }
        catch (Exception e)
        {
            LOG.error("RocketMQSource消费者消息异常", e);
            if (isPrintLog){
                System.out.println("RocketMQSource消费者消息异常:"+JSON.toJSONString(e));
            }
        }
    }
}
